#!/usr/bin/env python3
"""
生成三个不同模型的转录结果用于对比
"""

import json
from pathlib import Path
from faster_whisper import WhisperModel

# 测试视频
VIDEO_PATH = "/Volumes/Samsung_T5/数据爬取/B 站/关键词作品/狮子林/视频作品/BV13DGez8ETV苏州园林2-em class=keyword狮子林em 全网最详细游玩攻略.mp4"
OUTPUT_DIR = Path("output")
OUTPUT_DIR.mkdir(exist_ok=True)

# 模型配置
MODELS = ["tiny", "base", "small"]

def save_results(model_name, segments, info, video_path):
    """保存转录结果"""
    base_name = f"BV13DGez8ETV_{model_name}"

    # 收集文本
    all_segments = []
    full_text = []

    for segment in segments:
        seg = {
            "start": segment.start,
            "end": segment.end,
            "text": segment.text.strip()
        }
        all_segments.append(seg)
        full_text.append(segment.text.strip())

    # 保存 TXT
    txt_file = OUTPUT_DIR / f"{base_name}.txt"
    txt_file.write_text('\n'.join(full_text), encoding='utf-8')
    print(f"  ✓ {txt_file.name}")

    # 保存 SRT
    srt_file = OUTPUT_DIR / f"{base_name}.srt"
    with open(srt_file, 'w', encoding='utf-8') as f:
        for i, seg in enumerate(all_segments, 1):
            start = format_timestamp(seg['start'])
            end = format_timestamp(seg['end'])
            f.write(f"{i}\n")
            f.write(f"{start} --> {end}\n")
            f.write(f"{seg['text']}\n\n")
    print(f"  ✓ {srt_file.name}")

    # 保存 JSON
    json_file = OUTPUT_DIR / f"{base_name}.json"
    result_data = {
        "model": model_name,
        "video": str(video_path),
        "language": info.language,
        "duration": info.duration,
        "segments": all_segments,
        "full_text": '\n'.join(full_text)
    }
    json_file.write_text(
        json.dumps(result_data, ensure_ascii=False, indent=2),
        encoding='utf-8'
    )
    print(f"  ✓ {json_file.name}")

    return len(all_segments), len('\n'.join(full_text))

def format_timestamp(seconds):
    """格式化SRT时间戳"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

def main():
    print("=" * 80)
    print("生成三个模型的转录结果用于对比")
    print("=" * 80)
    print(f"\n视频: {Path(VIDEO_PATH).name[:60]}...")
    print(f"输出目录: {OUTPUT_DIR}\n")

    for model_name in MODELS:
        print(f"\n{'─' * 80}")
        print(f"模型: {model_name.upper()}")
        print(f"{'─' * 80}")

        # 加载模型
        print(f"加载模型...", end="", flush=True)
        model = WhisperModel(model_name, device="cpu", compute_type="int8")
        print(" 完成")

        # 转录
        print(f"转录中...", end="", flush=True)
        segments, info = model.transcribe(
            VIDEO_PATH,
            language="zh",
            initial_prompt="以下是普通话的句子。"
        )
        print(" 完成")

        # 保存结果
        print("保存文件:")
        seg_count, text_len = save_results(model_name, segments, info, VIDEO_PATH)

        print(f"\n统计:")
        print(f"  片段数: {seg_count}")
        print(f"  字符数: {text_len}")

    print("\n" + "=" * 80)
    print("✨ 所有转录完成！")
    print("=" * 80)
    print(f"\n输出文件位于: {OUTPUT_DIR.absolute()}")
    print("\n生成的文件:")
    for f in sorted(OUTPUT_DIR.glob("BV13DGez8ETV_*")):
        print(f"  - {f.name}")

if __name__ == '__main__':
    main()
